# %%
# from tkinter import NO
from calendar import EPOCH
from sched import scheduler
from turtle import pos
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from dataclasses import dataclass

import math

import torch.utils.data
# import pandas as pd
torch.manual_seed(seed=1024)
# %%
@dataclass
class GPTConfig:
    block_size: int = 512
    batch_size: int = 12
    n_layer: int = 6
    n_head: int = 12
    n_embd: int =768
    head_size: int = n_embd // n_head
    dropout: float = 0.1
    # ! 使用的是chatGpt-2的词表，大约有50257个token
    vocab_size: int = 50257
# %%
# * 模型结构
class SingleHeadAttention(nn.Module):
    """单头注意力机制

    Args:
        nn (_type_): _description_
    """
    def __init__(self, config) -> None:
        super().__init__()

        self.key = nn.Linear(in_features=config.n_embd, out_features=config.head_size)
        self.value = nn.Linear(in_features=config.n_embd, out_features=config.head_size)
        self.query = nn.Linear(in_features=config.n_embd, out_features=config.head_size)

        self.head_size = config.head_size

        # * nn.Module 提供的方法：用于早模块中注册一个不需要梯度计算的张量
        self.register_buffer(
            name='attention_mask',
            # * 创建下三角矩阵
            tensor=torch.tril(
                input=torch.ones(config.block_size, config.block_size)
            )
        )

        self.dropout = nn.Dropout(config.dropout)

    
    def forward(self, x):
        batch_size, seq_len, hidden_size = x.size()
        k = self.key(x)
        v = self.value(x)
        q = self.query(x)

        weight = q @ k.transpose(-2, -1)

        weight = weight.masked_fill(
            self.attention_mask[:seq_len, :seq_len] == 0,
            float('-inf')
        ) / math.sqrt(self.head_size)

        weight = F.softmax(input=weight, dim=-1)
        weight = self.dropout(weight)
        out = weight @ v
        return v
# %%
# * 多头注意力机制
class MultiHeadAttention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.heads = nn.ModuleList(
            SingleHeadAttention(config=config)
            for _ in range(config.n_head)
        )

        self.proj = nn.Linear(in_features=config.n_embd, out_features=config.n_embd)

        self.dropout = nn.Dropout(p=config.dropout)
    
    def forward(self, x):
        output = torch.cat(
            tensors=[h(x) for h in self.heads],
            dim = -1
        )
        output = self.proj(output)
        output = self.dropout(output)
        return output
# %%
# * Feed Network
class FeedForward(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(in_features=config.n_embd, out_features=4 * config.n_embd),
            nn.GELU(),
            nn.Linear(in_features=44 * config.n_embd, out_features=config.n_embd),
            nn.Dropout(p=config.dropout)
        )
    
    def forward(self, x):
        return self.net(x)
# %%
# ! 开始定义完整的block
class Block(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        head_size = config.n_embd // config.n_head
        self.attn = MultiHeadAttention(config=config)
        self.ffn = FeedForward(config=config)
        self.ln1 = nn.LayerNorm(normalized_shape=config.n_embd)
        self.ln2 = nn.LayerNorm(normalized_shape=config.n_embd)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.ffn(self.ln2(x))
        return x
# %%
# ! 构建完整的GPT
class GPT(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.token_enbdding_table = nn.Embedding(num_embeddings=config.vocab_size, embedding_dim=config.n_embd)
        self.position_embdding_table = nn.Embedding(num_embeddings=config.block_size, embedding_dim=config.n_embd)
        
        self.block = nn.Sequential(
            *[Block(config=config) for _ in range(config.n_layer)]
        )

        self.ln_final = nn.LayerNorm(normalized_shape=config.n_embd)
        self.lm_head = nn.Linear(in_features=config.n_embd, out_features=config.vocab_size, bias=False)

        # ? 怎么使用 疑惑 作用是什么？
        self.apply(fn=self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(tensor=module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(tensor=module.bias)
        
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(tensor=module.weight, mean=0.0, std=0.02)

    
    def forward(self, idx, targets=None):
        batch, seq_len = idx.size()
        token_emb = self.token_enbdding_table(idx)

        pos_emb = self.position_embdding_table(
            torch.arange(end=seq_len, device=idx.device)
        )

        x = token_emb + pos_emb
        x = self.block(x)
        x = self.ln_final(x)

        logits = self.lm_head(x)

        if targets is None:
            loss = None
        else:
            batch, seq_len, vocab_size = logits.size()
            logits = logits.view(batch * seq_len, vocab_size)
            targets = targets.view(batch * seq_len)
            loss = F.cross_entropy(input=logits, target=targets)
        
        return logits, loss

    def generate(self, idx, max_new_token):

        for _ in range(max_new_token):
            # * 序列太长，只取最后block_size个token
            idx_cond = idx if idx.size(1) <= self.block_size else idx[:, -self.block_size]

            logits, _ = self(idx_cond)
            # * 只关注最后时间步骤的预测 也可以理解成最后一个文本输入后的预测
            logits = logits[:, -1, :]

            probs = F.softmax(input=logits, dim=-1)
            # ? 采样下一个token 
            idx_next = torch.multinomial(input=probs, num_samples=1)

            idx = torch.cat(tensors=(idx, idx_next), dim=1)
        
        return idx
# %%
class MyDataset(Dataset):
    def __init__(self, path, block_size=512) -> None:
        super().__init__()
        import tiktoken
        self.enc = tiktoken.get_encoding(encoding_name="gpt2")
        self.block_size = block_size
        self.eos_token = self.enc.encode(
            "<|endoftext|>",
            allowed_special={"<|endoftext|>"}
        )[0]

        import json

        self.encoded_data = []

        self.max_lines = 1000

        raw_data = []

        with open(file=path, mode='r') as f:
            for i, line in enumerate(f):
                if i >= self.max_lines:
                    break
                try:
                    # ? 这里strip()函数的作用
                    text = json.loads(line.strip())['text']
                    raw_data.append(text)
                except json.JSONDecodeError:
                    continue
        full_encoded = []
        for text in raw_data:
            encoded_text = self.enc.encode(text)
            full_encoded.extend(encoded_text + [self.eos_token])
        
        # ! 将长文本分割成训练样本
        for i in range(0, len(full_encoded), self.block_size):
            # ? 多取一个Token作为目标
            chunk = full_encoded[i: i + self.block_size + 1]

            if len(chunk) < self.block_size + 1:
                chunk = chunk + [self.eos_token] * (self.block_size + 1 - len(chunk))
            
            self.encoded_data.append(chunk)

    def __len__(self) -> int:
        return len(self.encoded_data)
    
    def __getitem__(self, idx) -> tuple[torch.Tensor, torch.Tensor]:
        chunk = self.encoded_data[idx]
        # ? 为什么x不取最后一个值，y不取第一个值
        '''
            import numpy as np
            a = np.array([1,2,3,4,5,6,7,8,9,10])
            print(a[:-1]) # [1 2 3 4 5 6 7 8 9]
            print(a[1:]) # [ 2  3  4  5  6  7  8  9 10]

        '''
        x = torch.tensor(data=chunk[:-1], dtype=torch.long)
        y = torch.tensor(data=chunk[1:], dtype=torch.long)
        return x, y
    
    def encode(self, text) -> list[int]:
        return self.enc.encode(text=text)
    
    def decode(self, idx) -> str:
        return self.enc.decode(tokens=idx)
# %%
# 数据的格式
"""
{"text":"担任地点省市的区域运营中心的办理作业。承受总部相关KPI查核。\n1、了解新闻职业或媒体相关运营运营岗位，其间，应聘区域运营中心主任有3年以上当地干流媒体作业经验者优先，应聘事务主管有2年以上当地干流媒体作业经验者优先。\n2、交流才能强，抗压才能强，长于处理复杂情况，了解GR作业优先，能独立完结策划计划优先。具有独立开发客户才能。\n北京、天津、河北、山西、黑龙江、吉林、辽宁、上海、江苏、浙江、安徽、江西、福建、山东、河南、湖北、湖南、广东、海南、重庆、四川、贵州、云南、陕西等。"}
"""
# %%
path = ''
train_dataset = MyDataset(path)
train_dataset, val_dataset, = torch.utils.data.random_split(dataset=train_dataset, lengths=[0.9, 0.1])

train_loader = DataLoader(train_dataset, batch_size=12, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=12, shuffle=False)
# %%
model = GPT(GPTConfig())
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# * 打印模型参数

total_params = sum(p.numel() for p in model.parameters())
print(f"Total paramers:{total_params / 1e6} M")

optimizer = torch.optim.Adam(model.parameters(), lr = 3e-4)

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer, T_max=1000)
# %%
# todo 训练循环
def train(model, optimizer, scheduler, train_loader, val_loader, device, epoch):
    model.train()
    total_loss = 0

    for batch_idx, (x, y) in enumerate(train_loader):
        x, y = x.to(device), y.to(device)

        logits, loss = model(x, targets=y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()

        total_loss += loss.item()

        if batch_idx % 100 == 0:
            print(f'Epoch:{epoch}, Batch:{batch_idx}, Loss:{loss.item():.4f}')
    
    return total_loss
# todo 评估模型
def eval(model, val_loader, device):
    model.eval()
    val_loss=0
    with torch.no_grad():
        for x,y in val_loader:
            x,y = x.to(device), y.to(device)
            logits, loss = model(x, targets=y)
            val_loss += loss.item()
    return val_loss

for epoch in range(2):
    train_loss = train(model=model,optimizer=optimizer,scheduler=scheduler,train_loader=train_loader,val_loader=val_loader, device=device, epoch=epoch)
    val_loss = eval(model, val_loader, device)
    print(f'Epoch:{epoch}, Traub Loss:{train_loss/len(train_loader):.4f}, Val Loss:{val_loss/ len(val_loader):.4f}')

    avg_val_loss = val_loss / len(val_loader)

    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'val_loss': avg_val_loss
    }

    torch.save(checkpoint, f'checkpoint/model_epoch_{epoch}.pt')